import pygame
import random

# 初始化pygame
pygame.init()

# 设置屏幕大小
WIDTH, HEIGHT = 600, 400
CELL_SIZE = 20

# 颜色定义
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLACK = (0, 0, 0)

# 方向定义
UP = (0, -1)
DOWN = (0, 1)
LEFT = (-1, 0)
RIGHT = (1, 0)

# 初始化窗口
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("贪吃蛇")

# 定义蛇和食物
class Snake:
    def __init__(self):
        self.body = [(WIDTH // 2, HEIGHT // 2)]
        self.direction = RIGHT
        self.grow = False

    def move(self):
        head_x, head_y = self.body[0]
        new_head = (head_x + self.direction[0] * CELL_SIZE, head_y + self.direction[1] * CELL_SIZE)
        
        # 撞墙或撞自己时游戏结束
        if (new_head in self.body) or not (0 <= new_head[0] < WIDTH and 0 <= new_head[1] < HEIGHT):
            return False

        self.body.insert(0, new_head)
        if not self.grow:
            self.body.pop()
        self.grow = False
        return True

    def change_direction(self, new_direction):
        # 防止蛇掉头
        if (self.direction[0] + new_direction[0] != 0) or (self.direction[1] + new_direction[1] != 0):
            self.direction = new_direction

    def eat_food(self, food):
        if self.body[0] == food.position:
            self.grow = True
            return True
        return False

class Food:
    def __init__(self):
        self.position = (random.randint(0, (WIDTH - CELL_SIZE) // CELL_SIZE) * CELL_SIZE,
                         random.randint(0, (HEIGHT - CELL_SIZE) // CELL_SIZE) * CELL_SIZE)

    def spawn_new(self):
        self.position = (random.randint(0, (WIDTH - CELL_SIZE) // CELL_SIZE) * CELL_SIZE,
                         random.randint(0, (HEIGHT - CELL_SIZE) // CELL_SIZE) * CELL_SIZE)

# 主游戏循环
def main():
    clock = pygame.time.Clock()
    snake = Snake()
    food = Food()
    running = True

    while running:
        screen.fill(BLACK)

        # 处理事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP:
                    snake.change_direction(UP)
                elif event.key == pygame.K_DOWN:
                    snake.change_direction(DOWN)
                elif event.key == pygame.K_LEFT:
                    snake.change_direction(LEFT)
                elif event.key == pygame.K_RIGHT:
                    snake.change_direction(RIGHT)

        # 移动蛇
        if not snake.move():
            break  # 游戏结束

        # 吃食物
        if snake.eat_food(food):
            food.spawn_new()

        # 画蛇
        for segment in snake.body:
            pygame.draw.rect(screen, GREEN, pygame.Rect(segment[0], segment[1], CELL_SIZE, CELL_SIZE))

        # 画食物
        pygame.draw.rect(screen, RED, pygame.Rect(food.position[0], food.position[1], CELL_SIZE, CELL_SIZE))

        pygame.display.flip()
        clock.tick(10)  # 控制游戏速度

    pygame.quit()

if __name__ == "__main__":
    main()
